import json
import os
import argparse
import uuid
import time
from Crypto.Cipher import AES

import random
import string


textchars = bytearray({7,8,9,10,12,13,27} | set(range(0x20, 0x100)) - {0x7f})
is_binary_string = lambda bytes: bool(bytes.translate(None, textchars))

def is_binary(filename):
    with open(filename, 'rb') as f:
        return is_binary_string(f.read(1024))


key = '7smh3gXQk4rTuaI1cWVqfUxB2ZwLvy9G'
iv = '2DQiNTZBJj79oMqS'
signature = 'taichu'

block_size = 2048*16


def aes_en(data) -> bytes:
    if not isinstance(data, bytes):
        data = data.encode('utf-8')
    aes_obj = AES.new(key.encode("utf-8"), AES.MODE_CBC, iv.encode("utf-8"))
    aes_en_str = aes_obj.encrypt(data)
    return aes_en_str


def aes_de(data) -> bytes:
    if not isinstance(data, bytes):
        data = data.encode('utf-8')
    aes_de_obj = AES.new(key.encode("utf-8"), AES.MODE_CBC, iv.encode("utf-8"))
    aes_de_str = aes_de_obj.decrypt(data)

    return aes_de_str


def valid_one_file(root_dir, file_name):
    print('正在校验文件：', file_name)
    with open(os.path.join(root_dir, file_name), 'rb') as f:
        magic = f.read(len(signature))
        try:
            magic = magic.decode('utf-8')
        except UnicodeDecodeError:
            pass

        if magic == signature:
            print('文件签名check通过')

        try:
            header_raw = f.read(1024)
            file_header = aes_de(header_raw)
            file_header = json.loads(file_header)
            model_id = file_header.get('id', '')
            print('model_id 为', model_id)
        except Exception:
            print('获取模型ID异常！')


def enc_one_file(root_dir, file_name, info=None):
    print('正在加密文件：', file_name)
    begin = time.time()
    file_path = os.path.join(root_dir, file_name)
    # with open(file_path, 'rb') as f:
    #     magic = f.read(len(signature))
    #     try:
    #         magic = magic.decode('utf-8')
    #     except UnicodeDecodeError:
    #         pass

    #     if magic == signature:
    #         print('文件已加密过，跳过')
    #         return
    encrypt_tmp = os.path.join(root_dir, file_name + '_encrypt')
    with open(file_path, 'rb') as f:
        with open(encrypt_tmp, f'wb') as enc_f:
            # enc_f.write(signature.encode('utf-8'))
            # raw_h = padding_text(info, 1024)
            # enc_f.write(aes_en(raw_h.encode('utf-8')))

            while True:
                data = f.read(block_size)
                if len(data) < block_size:
                    enc_f.write(data)
                    break

                b = aes_en(data)
                enc_f.write(b)

    cost = time.time() - begin
    print(f'encrypt file: {file_name}, cost: {cost}s')
    os.remove(file_path)
    os.rename(encrypt_tmp, file_path)


def dec_one_file(root_dir, file_name):
    print('正在解密文件：', file_name)
    begin = time.time()
    file_path = os.path.join(root_dir, file_name)
    # with open(file_path, 'rb') as f:
    #     magic = f.read(len(signature)).decode('utf-8')
    #     if magic != 'taichu':
    #         print('文件签名不正确，已跳过')
    #         return
    decrypt_tmp = os.path.join(root_dir, file_name + '_decrypt')
    with open(file_path, 'rb') as f:
        with open(decrypt_tmp, 'wb') as dec_f:
            while True:
                data = f.read(block_size)
                if len(data) < block_size:
                    dec_f.write(data)
                    break

                b = aes_de(data)
                #print(f'encrypt_size: {block_size}, decrypt_size: {len(b)}')
                dec_f.write(b)

    cost = time.time() - begin
    print(f'decrypt file: {file_name}, cost: {cost}s')
    os.remove(file_path)
    os.rename(decrypt_tmp, file_path)


def padding_text(text, size=16):
    return text.ljust(size)  # 使用空字符 '\0' 进行填充


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--path', type=str, default=None, help='path to encrypt')
    parser.add_argument('--file', type=str, default=None, help='file to encrypt')
    parser.add_argument('--key', type=str, default=None, help='key to encrypt')
    parser.add_argument('--model_id', type=str, default=None, help='model id')

    parser.add_argument(choices=['enc', 'dec', 'valid'], dest='cmd', type=str, default='enc', help='encrypt or decrypt')
    args = parser.parse_args()

    # 生成随机32位密钥
    if args.key:
        key = args.key
    else:
        key = ''.join(random.sample(string.ascii_letters + string.digits, 32))
        print('random key:', key)

    if args.cmd == 'valid':
        if args.file:
            valid_one_file(os.path.dirname(args.file), os.path.basename(args.file))

        exit(0)

    if args.cmd == 'dec':
        if args.key is None:
            print('解密时必须指定密钥')
            exit(1)
        if args.file:
            dec_one_file(os.path.dirname(args.file), os.path.basename(args.file))

        if args.path:
            # 遍历文件夹
            for root, dirs, files in os.walk(args.path):
                for file in files:
                    if file.endswith('.pt') or file.endswith('.pth') or file.endswith('.bin'):
                        dec_one_file(root, file)
        exit(0)

    info = {
        'id': str(uuid.uuid4())
    }
    if args.model_id is not None:
        info['id'] = args.model_id

    if args.file:
        enc_one_file(os.path.dirname(args.file), os.path.basename(args.file), info=json.dumps(info, ensure_ascii=True))
        exit(0)

    if args.path:
        # 遍历文件夹
        begin = time.time()
        for root, dirs, files in os.walk(args.path):
            for file in files:
                if file.endswith('.bin'):
                    enc_one_file(root, file, info=json.dumps(info, ensure_ascii=True))
        cost = time.time() - begin
        print(f'encrypt path: {args.path}, total cost: {cost}s, blk_size: {block_size}bytes')
        exit(0)

# cp  /mnt/publish-data/train_data/llama/LLaMA13B/pytorch_model-00003-of-00006.bin
# python3 /mnt/publish-data/train_data/llama/encrypt_LLaMA13B/taichu_protect.py valid --file  /mnt/publish-data/train_data/llama/encrypt_LLaMA13B/pytorch_model-00001-of-00006.bin  --key OZzQ6lQi37lmL2kul4w3GTda1NLd9tdx

